
#include "KDTree.h"
#include "scene_parser.h"
#include <time.h>
#include <Windows.h>
#define SPLITNUMBERS 3

KDTree::KDTree(int maxDepth,size_t minTriangles)
{
	m_maxDepth = maxDepth;
	m_minTriangles = minTriangles;
	m_pRoot = new KDNode();
}

KDTree::~KDTree()
{
	for (int i = 0; i < m_sphereList.size(); ++i)
	{
		if ( m_sphereList[i] )
			delete m_sphereList[i];
	}
	m_sphereList.clear();
	
	// delete m_vpTriangles
	for (size_t i=0;i<m_pRoot->pvpTriangles->size();i++)
		delete (Triangle*)( (*(m_pRoot->pvpTriangles))[i] );

	// delete KDTree
	deleteKDNode(m_pRoot);
}

void KDTree::addTriangle(Triangle *pTriangle)
{
	m_pRoot->addTriangle(pTriangle);
}

void KDTree::addSphere(Sphere *pSphere)
{
	m_sphereList.push_back( pSphere );
}

void KDTree::deleteKDNode(KDNode *pNode)
{
	if ( pNode->pLeftNode )
		deleteKDNode( pNode->pLeftNode );
	if ( pNode->pRightNode )
		deleteKDNode( pNode->pRightNode );
	delete pNode;
}

void KDTree::buildTree()
{
	float startTime = timeGetTime();
	printf("\nstep 2: Kd-tree construction... ");
	srand(time(NULL));
	createNode(m_pRoot,0,2);
	printf("%.3fs\n",(timeGetTime()-startTime)/1000.f);
}

float SAH(KDNode *pLeft, KDNode *pRight/*,KDNode *pNode*/)
{
	float nnl = pLeft->pvpTriangles->size();
	float nnr = pRight->pvpTriangles->size();
	float vl = pLeft->aabb.getArea();
	float vr = pRight->aabb.getArea();
	//float v = pNode->aabb.getArea();
	return (nnl*vl + nnr*vr);///v;
}

inline void KDTree::trySplitValue(KDNode *&pLeftNode,KDNode *&pRightNode,KDNode *&pNode,
						  float &splitValue,float &minSAH, int xyz){
	KDNode *pLeft = new KDNode; KDNode *pRight = new KDNode;
	for (size_t i=0; i<pNode->pvpTriangles->size(); i++){
		Triangle *pTriangle = (*(pNode->pvpTriangles))[i];
		float minV = pTriangle->m_AABB.minXYZ[xyz];
		float maxV = pTriangle->m_AABB.maxXYZ[xyz];
		if ( maxV < splitValue )
			pLeft->addTriangle(pTriangle);
		else if ( minV > splitValue )
			pRight->addTriangle(pTriangle);
		else {
			pLeft->addTriangle(pTriangle);
			pRight->addTriangle(pTriangle);
		}
	}
	float sah = SAH(pLeft,pRight);
	if (sah < minSAH){
		minSAH = sah;
		pNode->splitValue = splitValue;
		if (pLeftNode) delete pLeftNode;
		if (pRightNode) delete pRightNode;
		pLeftNode = pLeft; pRightNode = pRight;
		pNode->splitXYZ = xyz;
	} else {
		delete pLeft; delete pRight;
	}
}
void KDTree::createNode(KDNode *pNode, int depth, int currXYZ)
{
	if (pNode->pvpTriangles->size()>m_minTriangles && depth<m_maxDepth){
		KDNode *pLeftNode = NULL;
		KDNode *pRightNode = NULL;

		float splitValue; float minSAH = FLT_MAX;
		if ( depth >= 20 ){
			for(int xyz=0; xyz <3; xyz++)
			for (int i=0;i<pNode->pvpTriangles->size();i++){
				splitValue = (*(pNode->pvpTriangles))[i]->m_AABB.minXYZ[xyz];
				trySplitValue(pLeftNode,pRightNode,pNode,splitValue,minSAH,xyz);
				splitValue = (*(pNode->pvpTriangles))[i]->m_AABB.maxXYZ[xyz];
				trySplitValue(pLeftNode,pRightNode,pNode,splitValue,minSAH,xyz);
			}
		} else {
			for(int xyz=0; xyz <3; xyz++){
				float maxValue = pNode->aabb.maxXYZ[xyz] - pNode->aabb.minXYZ[xyz];
				float midValue = (pNode->aabb.maxXYZ[xyz] + pNode->aabb.minXYZ[xyz]) / 2.f;
				for (int j=0; j < SPLITNUMBERS; j++){
					splitValue = midValue + (rand()%37777 - 37777.f/2.f)/37777.f * (maxValue);
					trySplitValue(pLeftNode,pRightNode,pNode,splitValue,minSAH,xyz);
				}
			}
		}

 		delete pNode->pvpTriangles;
 		pNode->pvpTriangles = NULL;
		createNode(pLeftNode,depth+1,pNode->splitXYZ);
		createNode(pRightNode,depth+1,pNode->splitXYZ);
		pNode->pLeftNode = pLeftNode;
		pNode->pRightNode = pRightNode;
	}
// 	else if (pNode->pvpTriangles->size() >= 50)
// 		printf("%d,%d ",depth,pNode->pvpTriangles->size());
}

bool KDTree::searchNode(const KDNode *pNode, const Ray *pRay, 
				Hit *pHit, float tmin, float tmax, int depth)
{
	if ( pNode->pvpTriangles!=NULL ){ // leaf
		if (pNode->pvpTriangles->size()==0) return false;
		if ( pHit )	{
			if( pNode->intersect(pRay,pHit,tmin,tmax) ){
				return true;}
			return false;
		}
		return pNode->intersectShadowRay(pRay,tmin,tmax);
	}

	Vec3f R = pRay->getOrigin(); Vec3f D = pRay->getDirection();
	if ( D[pNode->splitXYZ] > -FLT_EPSILON && D[pNode->splitXYZ] < FLT_EPSILON ){
		if ( R[pNode->splitXYZ] > pNode->splitValue )
			return searchNode(pNode->pRightNode,pRay,pHit,tmin,tmax,depth+1);
		else
			return searchNode(pNode->pLeftNode,pRay,pHit,tmin,tmax,depth+1);
	}

	KDNode *pNear,*pFar;
	if ( D[pNode->splitXYZ] >= FLT_EPSILON ){
		pNear = pNode->pLeftNode;
		pFar = pNode->pRightNode;
	} else {
		pNear = pNode->pRightNode;
		pFar = pNode->pLeftNode;
	}
	float thit = (pNode->splitValue-R[pNode->splitXYZ]) / D[pNode->splitXYZ];
	if ( thit >= tmax )
		return searchNode(pNear,pRay,pHit,tmin,tmax,depth+1);
	if ( thit <= tmin )
		return searchNode(pFar,pRay,pHit,tmin,tmax,depth+1);
	if ( searchNode(pNear,pRay,pHit,tmin,thit,depth+1) )
		return true;
	return searchNode(pFar,pRay,pHit,thit,tmax,depth+1);
}

bool KDTree::intersect(const Ray *pRay, Hit *pHit, float tmin, float tmax)
{
	bool hited = false;
	Hit hit;
	for (int i = 0; i < m_sphereList.size(); ++i)
	{
		if ( m_sphereList[i]->intersect(pRay, &hit, tmin, pHit->getT()) )
		{
			*pHit = hit;
			hited = true;
		}
	}
	if (searchNode(m_pRoot,pRay,&hit,tmin,pHit->getT(),0))
	{
		*pHit = hit;
		hited = true;
	}
	return hited;
}

bool KDTree::intersectShadowRay(const Ray *pRay, float tmin, float tmax)
{
	for (int i = 0; i < m_sphereList.size(); ++i)
	{
		if ( m_sphereList[i]->intersectShadowRay(pRay, tmin, tmax) )
			return true;
	}
 	return searchNode(m_pRoot,pRay,NULL,tmin,tmax,0);
}

void KDTree::debugMessage()
{
	printf("KDTree{\n}\n\n");
}